# data_loader.py
import os

import numpy as np
from scipy.stats import stats
from sklearn.model_selection import train_test_split

from auth_config import LOS_FP_PATH, LOS2_FP_PATH, NLOS_FP_PATH, \
    NLOS2_FP_PATH, MOV_FP_PATH, MOV2_FP_PATH, DL_FP_PATH, RAW_CSI_FEAT_DIM
from tools_1 import load_raw_csi, fuse_noise_with_attention, generate_universal_channel_noise, load_raw_csi_complex


# 在dataload_1.py中修改滑动窗口函数
def sliding_window(features, window_size, step, temporal_consistency=True):
    """
    增强版滑动窗口，添加时序一致性维护
    """
    num_samples, feat_dim = features.shape
    if num_samples < window_size:
        raise ValueError(f"样本数({num_samples})小于窗口大小({window_size})")

    num_windows = (num_samples - window_size) // step + 1
    if num_windows <= 0:
        raise ValueError(f"窗口参数不匹配")

    window_features = []
    for i in range(num_windows):
        start_idx = i * step
        end_idx = start_idx + window_size
        window = features[start_idx:end_idx, :]

        # 时序一致性维护
        if temporal_consistency and window_size > 1:
            # 计算窗口内时序差异
            diffs = np.abs(np.diff(window, axis=0)).mean(axis=0)
            # 基于时序稳定性动态加权
            weights = np.exp(-diffs / diffs.mean()) if diffs.mean() > 0 else np.ones_like(diffs)
            window = window * weights  # 时序稳定区域权重更高

        window_avg = window.mean(axis=0)
        window_features.append(window_avg)

    return np.array(window_features)


def filter_low_quality_samples(features, labels, threshold=0.3):
    """
    使用皮尔逊系数过滤低质量样本
    计算每个样本与同类样本均值的相关性，低于阈值的视为低质样本
    """
    if len(features) == 0:
        return features, labels

    filtered_features = []
    filtered_labels = []

    # 按类别计算均值
    unique_labels = np.unique(labels)
    for label in unique_labels:
        mask = (labels == label)
        class_features = features[mask]
        class_mean = np.mean(class_features, axis=0)

        # 计算每个样本与类别均值的皮尔逊系数
        for feat in class_features:
            corr, _ = stats.pearsonr(feat.flatten(), class_mean.flatten())
            if corr >= threshold:  # 保留高质量样本
                filtered_features.append(feat)
                filtered_labels.append(label)

    return np.array(filtered_features), np.array(filtered_labels)


# dataload_1.py（修正后）
def load_raw_train_data(
        device_list,
        source_scenario="LOS",
        use_dut11_noise=False,
        window_size=1,
        step=1,
        enable_filter=False,  # 新增：是否启用低质样本过滤
        quality_threshold=0.3,  # 新增：皮尔逊系数阈值
        enable_temporal_consistency=False  # 新增：是否启用时序一致性
):
    """加载目标设备单场景数据 + DUT11通用干扰，生成带多信道的训练集"""
    all_feats = []
    all_labels = []

    # 1. 提前生成DUT11的通用信道干扰（仅生成一次，复用）
    if use_dut11_noise:
        universal_noise = generate_universal_channel_noise()
        print(f"✅ 已加载DUT11通用信道干扰：{list(universal_noise.keys())}")

    # 2. 处理每个目标设备（DUT01~DUT10）：仅保留单层循环
    for idx, device in enumerate(device_list):  # 移除内层循环，只遍历一次设备列表
        print(f"正在处理设备 {device}（基础场景：{source_scenario}）")
        # 2.1 加载目标设备的单场景原始CSI
        target_csi = load_raw_csi_complex(
            device_name=device,
            scenario=source_scenario,
        )

        # 2.2 融入DUT11的多信道干扰
        if use_dut11_noise:
            target_csi = fuse_noise_with_attention(target_csi, universal_noise)
            print(f"融合后设备{device}数据形状：{target_csi.shape}")

        # 2.3 转换为幅度+相位双通道
        amplitude = np.abs(target_csi)
        phase = np.angle(target_csi)
        csi_feat = np.stack([amplitude, phase], axis=1)  # [样本数, 2, 子载波数]

        # 2.4 滑动窗口叠帧
        flat_feat = csi_feat.reshape(csi_feat.shape[0], -1)
        windowed_feat = sliding_window(flat_feat, window_size, step, temporal_consistency=enable_temporal_consistency)

        # 2.5 生成标签（与当前设备索引匹配）
        window_labels = np.full(shape=(windowed_feat.shape[0],), fill_value=idx, dtype=int)

        all_feats.append(windowed_feat)
        all_labels.append(window_labels)

    # 3. 合并所有设备数据
    X_train = np.concatenate(all_feats, axis=0)
    y_train = np.concatenate(all_labels, axis=0)

    # 新增：低质样本裁剪（如果启用）
    if enable_filter:
        X_train, y_train = filter_low_quality_samples(X_train, y_train, threshold=quality_threshold)
        print(f"✅ 低质样本裁剪完成：保留 {X_train.shape[0]} 个样本（阈值={quality_threshold}）")

    # 恢复为2通道格式
    X_train = X_train.reshape(-1, 2, RAW_CSI_FEAT_DIM)  # 使用配置中的维度，避免硬编码
    print(f"\n✅ 最终训练集：{X_train.shape[0]}个样本（含DUT11多信道干扰），特征维度{X_train.shape[1:]}")
    print(f"✅ 标签样本数：{y_train.shape[0]}")  # 新增检查，确保与X_train样本数一致
    return X_train, y_train

# dataload_1.py（修改后）
# 可选：测试集加载函数也可添加时序一致性参数（按需修改）
def load_raw_test_data(
        device_list,
        target_scenario="NLOS",
        window_size=3,
        step=1,
        enable_temporal_consistency=False  # 新增
):
    """加载原始CSI测试数据（带时序一致性处理）"""
    all_feats = []
    all_labels = []
    all_device_names = []

    for idx, device in enumerate(device_list):
        csi_feat = load_raw_csi(device, scenario=target_scenario)
        flat_feat = csi_feat.reshape(csi_feat.shape[0], -1)
        # 应用时序一致性
        windowed_feat = sliding_window(
            flat_feat,
            window_size,
            step,
            temporal_consistency=enable_temporal_consistency
        )
        all_feats.append(windowed_feat)

        window_labels = np.full(shape=(windowed_feat.shape[0],), fill_value=idx, dtype=int)
        all_labels.append(window_labels)
        all_device_names.extend([device] * windowed_feat.shape[0])

    X_test = np.concatenate(all_feats, axis=0)
    X_test = X_test.reshape(-1, 2, 50)
    y_test = np.concatenate(all_labels, axis=0)
    print(f"✅ 原始CSI测试集（叠帧后）：{X_test.shape[0]}个样本，特征维度{X_test.shape[1:]}")
    return X_test, y_test, all_device_names